import numpy as np
import matplotlib.pyplot as plt
from pykrige.ok import OrdinaryKriging

# 生成虚拟数据集
np.random.seed(42)
num_samples = 30  # 原始观测点的数量
x_observed = np.random.uniform(0, 100, num_samples)
y_observed = np.random.uniform(0, 100, num_samples)
z_observed = np.sin(x_observed / 10) * np.cos(y_observed / 10) + np.random.normal(0, 0.1, num_samples)

# 定义插值网格
grid_x = np.linspace(0, 100, 200)
grid_y = np.linspace(0, 100, 200)
grid_x, grid_y = np.meshgrid(grid_x, grid_y)

# 使用普通克里金插值 (Ordinary Kriging)
OK = OrdinaryKriging(
    x_observed, y_observed, z_observed,
    variogram_model="spherical",
    verbose=False, enable_plotting=False
)
z_kriging, z_variance = OK.execute("grid", grid_x[0], grid_y[:, 0])

# 创建图形
fig, ax = plt.subplots(1, 3, figsize=(18, 6))

# 原始数据分布图
scatter = ax[0].scatter(x_observed, y_observed, c=z_observed, cmap="coolwarm", s=50, edgecolor='k')
ax[0].set_title("Observed Data Distribution")
ax[0].set_xlabel("X coordinate")
ax[0].set_ylabel("Y coordinate")
fig.colorbar(scatter, ax=ax[0], label="Observed Value")

# 克里金插值结果图
im = ax[1].imshow(z_kriging, extent=(0, 100, 0, 100), origin="lower", cmap="coolwarm")
ax[1].set_title("Kriging Interpolation Result")
ax[1].set_xlabel("X coordinate")
ax[1].set_ylabel("Y coordinate")
fig.colorbar(im, ax=ax[1], label="Interpolated Value")

# 插值误差（方差）分布图
im_var = ax[2].imshow(z_variance, extent=(0, 100, 0, 100), origin="lower", cmap="YlOrBr")
ax[2].set_title("Interpolation Variance")
ax[2].set_xlabel("X coordinate")
ax[2].set_ylabel("Y coordinate")
fig.colorbar(im_var, ax=ax[2], label="Variance")

# 显示所有图形
plt.tight_layout()
plt.show()